from core import consts
from core.exprs import Expr, Literal, Variable, Application, Body, Operator, IF, DefineVar, DefineFunction
from core.mtoken import Token, TokenType
from core.printer import ASTPrinter


class Parser:

    def __init__(self):
        self.current = 0

    def _reset(self, tokens: list[Token]):
        self.tokens = tokens
        self.current = 0

    def parse(self, tokens: list[Token]) -> list[Expr]:
        self._reset(tokens)
        forms = []
        while not self.check(TokenType.EOF):
            form = self.form()
            forms.append(form)
        self.debug(forms)
        return forms

    def form(self) -> Expr:
        if not self.check(TokenType.LEFT_PAREN):
            return self.expression()
        else:
            self.match(TokenType.LEFT_PAREN, "Expected '('")
            if self.match(TokenType.DEFINE):
                return self.define()
            else:
                return self.expression(True)

    def define(self):
        variables = []
        defineVar = True
        if self.match(TokenType.LEFT_PAREN):
            defineVar = False
            variables.append(self.variable())
            while not self.match(TokenType.RIGHT_PAREN):
                variables.append(self.variable())
        else:
            variables.append(self.variable())

        body = self.body()
        self.consume(TokenType.RIGHT_PAREN, "define Expected ')'")
        if defineVar:
            return DefineVar(variables[0].name, body)
        else:
            return DefineFunction(variables[0].name, variables[1:], body)

    def body(self):
        defines = []
        while self.check(TokenType.LEFT_PAREN) and self.checkNext(TokenType.LEFT_PAREN):
            self.consume(TokenType.LEFT_PAREN, "Expected '('")
            self.consume(TokenType.DEFINE, "Expected 'define'")
            defines.append(self.define())

        expressions = []
        first = self.expression()
        expressions.append(first)
        while not self.check(TokenType.RIGHT_PAREN):
            expressions.append(self.expression())
        if len(expressions) == 0:
            raise Exception("Expected expression")
        else:
            return Body(defines, expressions)

    def variable(self):
        if self.match(TokenType.IDENTIFIER):
            return Variable(self.previous())
        else:
            raise Exception("Expected variable")

    def advance(self):
        self.current += 1

    def consume(self, ttype, message):
        if self.check(ttype):
            self.advance()
        else:
            raise Exception(message)

    def checkNext(self, ttype: TokenType):
        if self.current + 1 < len(self.tokens):
            return self.tokens[self.current + 1].ttype == ttype
        else:
            return False

    def check(self, ttype: TokenType) -> bool:
        if self.current < len(self.tokens):
            return self.tokens[self.current].ttype == ttype
        else:
            if ttype == TokenType.EOF:
                return True
        return False

    def match(self, *ttypes: TokenType):
        for ttype in ttypes:
            if (self.check(ttype)):
                self.advance()
            else:
                return False
        return True

    def IF(self) -> Expr:
        condition = self.expression()
        consequence = self.expression()
        alternative = None

        if not self.match(TokenType.RIGHT_BRACE):
            alternative = self.expression()
        return IF(condition, consequence, alternative)

    def expression(self, match_left_paren=False) -> Expr:
        if match_left_paren or self.match(TokenType.LEFT_PAREN):
            if self.match(TokenType.IF):
                return self.IF()
            else:
                return self.application()
        else:
            if self.match(TokenType.NUMBER):
                first = Literal(self.previous().value)
            elif self.match(TokenType.STRING):
                first = Literal(self.previous().value)
            elif self.match(TokenType.TRUE):
                first = Literal(True)
            elif self.match(TokenType.FALSE):
                first = Literal(False)
            elif self.match(TokenType.PLUS):
                first = Operator(self.previous())
            elif self.match(TokenType.MINUS):
                first = Operator(self.previous())
            elif self.match(TokenType.STAR):
                first = Operator(self.previous())
            elif self.match(TokenType.SLASH):
                first = Operator(self.previous())
            elif self.match(TokenType.LESS):
                first = Operator(self.previous())
            elif self.match(TokenType.GREATER):
                first = Operator(self.previous())
            elif self.match(TokenType.EQUAL_EQUAL):
                first = Operator(self.previous())
            elif self.match(TokenType.LESS_EQUAL):
                first = Operator(self.previous())
            elif self.match(TokenType.GREATER_EQUAL):
                first = Operator(self.previous())
            elif self.match(TokenType.BANG_EQUAL):
                first = Operator(self.previous())
            elif self.match(TokenType.BANG):
                first = Operator(self.previous())
            elif self.match(TokenType.AND):
                first = Operator(self.previous())
            elif self.match(TokenType.OR):
                first = Operator(self.previous())
            elif self.match(TokenType.IDENTIFIER):
                first = Variable(self.previous())
            elif self.check(TokenType.LEFT_PAREN):
                first = self.expression()
            else:
                raise Exception("Unsupported expression")
            return first

    def previous(self):
        return self.tokens[self.current - 1]

    def peek(self) -> Token:
        return self.tokens[self.current]

    def application(self):
        sequences = []
        while not self.check(TokenType.RIGHT_PAREN) and not self.check(TokenType.EOF):
            sequences.append(self.expression())

        if len(sequences) == 0:
            raise Exception("Expected expression")
        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after sequence")
        return Application(sequences[0], sequences[1:])

    def debug(self, exprs: list[Expr]):
        if consts.PARSE:
            printer = ASTPrinter()
            print("================= ast =================")
            for expr in exprs:
                printer.print(expr)
            print("========================================")
